#include <string.h>
#include <stdlib.h>
#include "avl.h"

#define MAX(X, Y) (X > Y ? X : Y)

avl_t *avl_create(int key_limit)
{
    avl_t *avl = (avl_t *) malloc(sizeof(avl_t));

    avl->size = 0;
    avl->root = NULL;
    avl->key_limit = key_limit;

    return avl;
}

static void avl_clear_recursive(node_t *node)
{
    if (node == NULL)
        return;

    avl_clear_recursive(node->left);
    avl_clear_recursive(node->right);

    free(node);
}

void avl_destroy(avl_t *avl)
{
    avl_clear_recursive(avl->root);

    free(avl);
}

static int height(node_t *node)
{
    return node == NULL ? -1 : node->height;
}

/* Function to balance de tree. */
static node_t *single_rotate_with_left(node_t *n2)
{
    node_t *n1;

    n1 = n2->left;
    n2->left = n1->right;
    n1->right = n2;

    n2->height = MAX(height(n2->left), height(n2->right)) + 1;
    n1->height = MAX(height(n1->left), n2->height) + 1;

    return n1;
}

/* Function to balance de tree. */
static node_t *single_rotate_with_right(node_t *n1)
{
    node_t *n2;

    n2 = n1->right;
    n1->right = n2->left;
    n2->left = n1;

    n1->height = MAX(height(n1->left), height(n1->right)) + 1;
    n2->height = MAX(height(n2->right), n1->height) + 1;

    return n2;
}

/* Function to balance de tree. */
static node_t *double_rotate_with_left(node_t *n3)
{
    n3->left = single_rotate_with_right(n3->left);

    return single_rotate_with_left(n3);
}

/* Function to balance de tree. */
static node_t *double_rotate_with_right(node_t *n1)
{
    n1->right = single_rotate_with_left(n1->right);

    return single_rotate_with_right(n1);
}

/* Function to safely copy the key. */
static void
copy_key(char *dest, const char *key, int key_limit)
{
    int i;
    const char *p;

    for (p = key, i = 0; *p != '\0' && i < key_limit; ++p, ++i)
        dest[i] = *p;
    for (; i < key_limit; ++i)
        dest[i] = '_';
    dest[i] = '\0';
}

/* Function to safely compare keys. */
static int
cstrcmp(const char *a, const char *b, int key_limit)
{
    int i;
    for (i = 0; i < key_limit; i++)
    {
        /* '\0' and '_' must be treat equally. */
        if (a[i] == '\0' || a[i] == '_')
            return b[i] == '_' ? 0 : - b[i];

        if (b[i] == '\0' || b[i] == '_')
            return a[i] == '_' ? 0 : a[i];

        if (a[i] != b[i])
            return a[i] - b[i];
    }

    return 0;
}

static node_t *
avl_insert_recursive(node_t *node, const char *key, data_t data, node_t **new_node, int key_limit, int *ret)
{
    /* Place to store the new node. */
    if (node == NULL)
    {
        node = (node_t *) malloc(sizeof(node_t));

        copy_key(node->key, key, key_limit);
        node->data = data;
        node->height = 0;
        node->left = node->right = NULL;
        node->removed = 0;

        *new_node = node;
        *ret = 1;
    }
    /* The new word is lexicographically smaller than the other. */
    else if (cstrcmp(key, node->key, key_limit) < 0)
    {
        node->left = avl_insert_recursive(node->left, key, data, new_node, key_limit, ret);
        /* Balance if necessary. */
        if (height(node->left) - height(node->right) == 2)
        {
            if (cstrcmp(key, node->left->key, key_limit) < 0)
                node = single_rotate_with_left(node);
            else
                node = double_rotate_with_left(node);
        }
    }
    /* The new word is lexicographically greater than the other. */
    else if (cstrcmp(key, node->key, key_limit) > 0)
    {
        node->right = avl_insert_recursive(node->right, key, data, new_node, key_limit, ret);
        /* Balance if necessary. */
        if (height(node->right) - height(node->left) == 2)
        {
            if (cstrcmp(key, node->right->key, key_limit) > 0)
                node = single_rotate_with_right(node);
            else
                node = double_rotate_with_right(node);
        }
    }
    else
    {
        /* The key already exists. Return it. */
        *new_node = node;
    }

    /* Update height. */
    node->height = MAX(height(node->left), height(node->right)) + 1;

    return node;
}

node_t *avl_insert(avl_t *avl, const char *key, data_t data)
{
    int ret = 0;
    node_t *new_node = NULL;

    avl->root = avl_insert_recursive(avl->root, key, data, &new_node, avl->key_limit, &ret);
    if(ret)
        avl->size++;

    return new_node;
}

static node_t *avl_search_recursive(node_t *node, const char *key, int key_limit)
{
    if (node == NULL)
        return NULL;

    int cmp = cstrcmp(key, node->key, key_limit);
    if (cmp == 0)
    {
        /* The node was removed. */
        if (node->removed)
            return NULL;
        else
            return node;
    }
    if (cmp < 0)
        return avl_search_recursive(node->left, key, key_limit);
    return avl_search_recursive(node->right, key, key_limit);
}

node_t *avl_search(avl_t *avl, const char *key)
{
    return avl_search_recursive(avl->root, key, avl->key_limit);
}

static node_t *
avl_remove_recursive(node_t *node, const char *key, int *ret, int key_limit)
{
    if (node == NULL)
        return NULL;

    /* The key is lexicographically smaller than the other. */
    if (cstrcmp(key, node->key, key_limit) < 0)
    {
        node->left = avl_remove_recursive(node->left, key, ret, key_limit);
    }
    /* The key is lexicographically greater than the other. */
    else if (cstrcmp(key, node->key, key_limit) > 0)
    {
        node->right = avl_remove_recursive(node->right, key, ret, key_limit);
    }
    /* Remove the node. In fact, only mark it. */
    else if (node->removed == 0)
    {
        *ret = 1;
        node->removed = 1;
    }

    return node;
}

void avl_remove(avl_t *avl, const char *key)
{
    int ret = 0;

    avl->root = avl_remove_recursive(avl->root, key, &ret, avl->key_limit);

    if (ret)
        avl->size--;
}


